import numpy as np


class SRI(object):
    def __init__(self, N=51):
        '''
        :param N(int): 养护路段总长度(Km)，向上取整
        '''

        super().__init__()
        # 路面横向力系数预测模型参数表
        self.a = np.array([[62.748, 62.544, 63.278, 62.889],
                           [68.455, 68.229, 68.563, 67.996],
                           [69.883, 69.269, 70.658, 69.295],
                           [71.679, 70.212, 72.883, 70.541],
                           [71.226, 70.443, 72.986, 69.136],
                           [72.118, 70.469, 72.314, 70.417],
                           [72.235, 70.883, 70.348, 70.134],
                           [72.554, 70.913, 70.459, 70.213]])
        self.b = np.array([[-0.047, -0.053, -0.061, -0.066],
                           [-0.043, -0.049, -0.060, -0.064],
                           [-0.041, -0.046, -0.059, -0.059],
                           [-0.041, -0.046, -0.059, -0.059],
                           [-0.041, -0.043, -0.058, -0.056],
                           [-0.041, -0.042, -0.055, -0.055],
                           [-0.040, -0.042, -0.051, -0.055],
                           [-0.040, -0.042, -0.050, -0.054]])
        self.SRI_min = np.array([35.0])

        self.N = N

    def SFC(self, t, j, b):
        '''
        :param t: 养护年
        :param j: 养护类别
        :param b: 路段交通量
        :return:
        '''

        x1 = self.a * np.exp(self.b * t)
        res = np.stack([x1[j, b] for _ in range(self.N)]).squeeze()
        return res

    def sfc_new(self, t, js, b):
        '''
        :param t: 养护年
        :param j: 养护类别
        :param b: 路段交通量
        :return:
        '''

        x1 = self.a * np.exp(self.b * t)
        res =np.stack([x1[:, b][jx] for jx in js]).squeeze()
        return res

        # res = np.stack([x1[j, b] for _ in range(self.N)]).squeeze()
        # return res

    def SRI(self, sfc):
        return (100.0 - self.SRI_min) / (1.0 + 28.6 * np.exp(-0.105 * sfc)) + self.SRI_min


if __name__ == '__main__':
    T = 5
    sfc = SRI()
    j = np.random.randint(0, 8, (10, 1))
    b = np.random.randint(0, 4, (10, 1))

    for t in range(T):
        xx = sfc.SFC(t, j, b)
        print(xx)
        sri = sfc.SRI(xx)
        print(sri)
